# test 
import os
import json
from face_search import get_config
from face_search import SearchEngine


def build_facedb():
    cfg = get_config('./configs/face_search_v0.2.0.yaml')
    engine = SearchEngine(cfg)

    facedb_folder = '../face_search_data/facedb_v1'
    for image_name in os.listdir(facedb_folder):
        image_path = os.path.abspath(os.path.join(facedb_folder, image_name))
        person_name = image_name.split('.')[0]
        print('adding {}: {}'.format(person_name, image_path))
        engine.add_face(image_path, person_name)
    engine.save()


def search_face(face_folder, result_folder):
    if not os.path.exists(result_folder):
        os.makedirs(result_folder)
    
    cfg = get_config('./configs/face_search_v0.2.0.yaml')
    engine = SearchEngine(cfg)

    # facetest_folder = '../face_search_data/facetest_v1'
    for image_name in os.listdir(face_folder):
        print(image_name)
        filename = image_name.split('.')[0]
        json_path = os.path.join(result_folder, '{}.json'.format(filename))
        if os.path.exists(json_path):
            continue
        image_path = os.path.abspath(os.path.join(face_folder, image_name))
        face_info = engine.search_by_image(image_path, topk=3)
        print(face_info)
        with open(json_path, 'w', encoding='utf-8') as f:
            data = json.dumps(face_info)
            f.write(data)


def load_result_json(json_path):
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.loads(f.read())
    return data


if __name__ == '__main__':
    # build_facedb()

    face_folder = '../face_search_data/facetest_v1'
    result_folder = '../face_search_data/facetest_v1_result2.0'
    # data = load_result_json(os.path.join(result_folder, '丁仲礼_0.json'))
    # print(data)
    search_face(face_folder, result_folder)
